package io.kiki.stack.http.feign.SpringMvcContract.loadbalancer;

import io.kiki.stack.http.feign.Request;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.LoadBalancerProperties;
import org.springframework.cloud.client.loadbalancer.reactive.ReactiveLoadBalancer;

import java.net.URI;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

/**
 * To add X-Forwarded-Host and X-Forwarded-Proto Headers.
 */
public class XForwardedHeadersTransformer implements LoadBalancerFeignRequestTransformer {

    private final ReactiveLoadBalancer.Factory<ServiceInstance> factory;

    public XForwardedHeadersTransformer(ReactiveLoadBalancer.Factory<ServiceInstance> factory) {
        this.factory = factory;
    }

    @Override
    public Request transformRequest(Request request, ServiceInstance serviceInstance) {
        if (serviceInstance == null) {
            return request;
        }
        LoadBalancerProperties.XForwarded xForwarded = factory.getProperties(serviceInstance.getServiceId()).getXForwarded();
        if (xForwarded.isEnabled()) {
            Map<String, Collection<String>> headers = new HashMap<>(request.headers());
            URI uri = URI.create(request.url());
            String xForwardedHost = uri.getHost();
            String xForwardedProto = uri.getScheme();
            headers.put("X-Forwarded-Host", Collections.singleton(xForwardedHost));
            headers.put("X-Forwarded-Proto", Collections.singleton(xForwardedProto));
            request = Request.create(request.httpMethod(), request.url(), headers, request.body(), request.charset(), request.requestTemplate());
        }
        return request;
    }

}
